"""
上下文基类，回测过程中净值的统计，一些信息的更新，以及结果测输出由上下文基类控制
"""
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt

from dytz.bar import *

__all__ = ['Context', ]


class Context:
    """上下文基类， 负责结果的统计，回测过程中一些信息的更新，和结果的输出"""

    max_high = None
    max_low = None
    min_high = None
    min_low = None

    entry_price = None

    def __init__(self, cash: float = 1000000.0, slippage: float = 0.001,
                 multiplier: int = 10, charge_ratio: float = 0.5 / 10000):
        """
        构造器
        :param cash: 本金
        :param slippage: 滑点
        :param multiplier: 合约乘数，为一手要购买的量
        :param charge_ratio: 手续费
        """
        self.cash = self.init_cash = cash    # 资金账户
        self.cap_count = 0                  # 资产账户
        self.slippage_r = slippage
        self.multiplier = multiplier
        self.charge = 0
        self.charge_ratio = charge_ratio
        self.charge_all = []
        self.pos = 0
        self.rnt_ = 0
        self.rnt_all = []
        self.rnt = []
        self.date = []
        self.inf = []
        self.close = []
        self.cash_all = []
        self.pri = []
        self.tar = []
        self.slippage = 0
        self.slippage_all = []
        self.cap_all = []
        self.target_pos = None
        self.target_price = None

    def on_bar(self, bar: Bar) -> None:
        """
        每次bar线生成后调用
        :param bar: Bar线
        :return: None
        """
        if self.pos != 0:
            if self.min_low is None:
                self.min_low = bar.LOW
                self.max_low = bar.LOW
                self.min_high = bar.HIGH
                self.max_high = bar.HIGH
            else:
                self.max_low = max(bar.LOW, self.max_low)
                self.min_low = min(bar.LOW, self.min_low)
                self.max_high = max(bar.HIGH, self.max_high)
                self.min_high = min(bar.HIGH, self.min_high)

        if self.target_pos is not None:  #  仓位控制
            self.__open_long(bar.AVGPRICE)
            self.__open_short(bar.AVGPRICE)
            self.__cover_long(bar.AVGPRICE)
            self.__cover_short(bar.AVGPRICE)


        self.target_pos = None


    def on_sta(self, bar: Bar) -> None:
        """
        收益率的统计, 统计时间为每天盘后
        :param bar: Bar线
        :return: None
        """
        self.date.append(bar.date)  # 日期统计
        if self.pos != 0:  # 获利统计
            self.rnt_ = (bar.CLOSE - self.entry_price) * self.pos
            self.entry_price = bar.CLOSE
            self.cap_count += self.rnt_
        else:
            pass
        self.tar.append(self.target_pos)
        self.rnt.append(self.rnt_)
        self.cash_all.append(self.cash)
        self.cap_all.append(self.cap_count)
        self.inf.append(self.pos)
        self.close.append(bar.CLOSE)
        self.pri.append(bar.AVGPRICE)
        self.charge_all.append(self.charge)
        self.slippage_all.append(self.slippage)
        self.rnt_ = 0
        self.charge = 0
        self.slippage = 0

    def __open_long(self, price: float) -> None:
        """
        开多函数
        :param price: 价格
        :return: None
        """
        if self.target_pos > 0 and self.pos == 0:  # 当没有持仓时
            pos = (self.cash / (price * (1 + self.slippage_r + self.charge_ratio)) // self.multiplier) * self.multiplier
            self.pos += pos
            self.entry_price = price
            self.cap_count = pos * price
            self.charge = self.cap_count * self.charge_ratio
            self.slippage = self.cap_count * self.slippage_r
            self.cash -= self.cap_count + self.charge + self.slippage
        elif self.target_pos > 0 and self.pos < 0:  # 当反方向持仓时
            self.__cover_short(price)
            pos = (self.cash / (price * (1 + self.slippage_r)) // self.multiplier) * self.multiplier
            self.pos += pos
            self.entry_price = price
            self.cap_count = pos * price
            self.charge = self.cap_count * self.charge_ratio
            self.slippage = self.cap_count * self.slippage_r
            self.cash -= self.cap_count + self.charge + self.slippage
        else:
            pass

    def __open_short(self, price: float) -> None:
        """
        开多函数
        :param price: 价格
        :return: None
        """
        if self.target_pos < 0 and self.pos == 0:
            pos = (self.cash / (price * (1 + self.slippage_r + self.charge_ratio)) // self.multiplier) * self.multiplier
            self.pos -= pos
            self.entry_price = price
            self.cap_count = pos * price
            self.charge = self.cap_count * self.charge_ratio
            self.slippage = self.cap_count * self.slippage_r
            self.cash -= self.cap_count + self.charge + self.slippage
        elif self.target_pos < 0 and self.pos > 0:
            self.__cover_long(price)
            pos = (self.cash / (price * (1 + self.slippage_r)) // self.multiplier) * self.multiplier
            self.pos -= pos
            self.entry_price = price
            self.cap_count = pos * price
            self.charge = self.cap_count * self.charge_ratio
            self.slippage = self.cap_count * self.slippage_r
            self.cash -= self.cap_count + self.charge + self.slippage
        else:
            pass

    def __cover_long(self, price: float) -> None:
        """
        平多函数
        :param price: 价格
        :return: None
        """
        if self.pos > 0 and self.target_pos == 0:
            self.rnt_ = (price - self.entry_price) * self.pos
            self.cap_count += self.rnt_
            self.charge = self.cap_count * self.charge_ratio
            self.slippage = self.cap_count * self.slippage_r
            self.cash -= self.charge + self.slippage
            self.cash += self.cap_count
            self.cap_count = 0
            self.max_low = None
            self.max_high = None
            self.min_high = None
            self.min_low = None
            self.pos = 0

    def __cover_short(self, price: float) -> None:
        """
        平空函数
        :param price: 价格
        :return: None
        """
        if self.pos < 0 and self.target_pos == 0:
            self.rnt_ = (price - self.entry_price) * self.pos
            self.cap_count += self.rnt_
            self.charge = self.cap_count * self.charge_ratio
            self.slippage = self.cap_count * self.slippage_r
            self.cash -= self.charge + self.slippage
            self.cash += self.cap_count
            self.cap_count = 0
            self.max_low = None
            self.max_high = None
            self.min_high = None
            self.min_low = None
            self.pos = 0

    def open_long(self, price: float = None) -> None:
        """
        开多函数
        :param price: 价格
        :return: None
        """
        self.target_pos = 1

    def open_short(self, price: float = None) -> None:
        """
        开多函数
        :param price: 价格
        :return: None
        """
        self.target_pos = -1

    def cover_long(self, price: float = None) -> None:
        """
        平多函数
        :param price: 价格
        :return: None
        """
        self.target_pos = 0

    def cover_short(self, price = None):
        """
        平空函数
        :param price: 价格
        :return: None
        """
        self.target_pos = 0

    def table(self, file_name: str) -> None:
        """
        输出净值表
        :param file_name: 文件名
        :return: None
        """
        self.res = pd.DataFrame(index=self.date)
        self.res.loc[:, 'cash'] = self.cash_all
        self.res.loc[:, 'price'] = self.pri
        self.res.loc[:, 'close'] = self.close
        self.res.loc[:, 'position'] = self.inf
        self.res.loc[:, 'charge'] = self.charge_all
        self.res.loc[:, 'slippage'] = self.slippage_all
        self.res.loc[:, 'rnt'] = np.asarray(self.rnt)
        self.res.loc[:, 'pos_cap'] = np.asarray(self.cap_all)
        self.res.loc[:, 'cap'] = np.asarray(self.cash_all) + np.asarray(self.cap_all)
        self.res.loc[:, 'rnt_ratio'] = ((self.res.loc[:, 'cap'] - self.res.loc[:, 'cap'].shift(1)) /
                                        self.res.loc[:, 'cap'].shift(1))
        k_ = self.res.loc[:, 'rnt_ratio'] + 1
        self.res.loc[:, 'line_rnt'] = (np.cumprod(k_))
        self.avg = self.res.loc[:, 'line_rnt'][-1] ** (250 / len(self.res)) - 1
        self.std = np.std(self.res.loc[:, 'rnt_ratio']) * 252 ** 0.5
        self.shape = (self.avg - 0.03) / self.std
        maxdd = (self.res.loc[:, 'cap'].cummax() - self.res.loc[:, 'cap']) / self.res.loc[:, 'cap'].cummax()
        self.res.loc[:, 'cummaxdd'] = maxdd
        self.res.loc[:, 'target'] = np.asarray(self.tar)
        self.maxDD = maxdd.max()
        print(f"avg     {self.avg}\nstd       {self.std}\nshape       " +
              f" {self.shape}\nmaxdd      {self.maxDD}\ncalmar        {self.avg / self.maxDD}")
        self.calmar = (self.avg - 0.03) / self.maxDD
        self.back_res = {
            'net': self.res.loc[:, 'cap'][-1] / self.init_cash,
            'avg': self.avg,
            'std': self.std,
            'shape': self.shape,
            'maxDD': self.maxDD,
            'calmar': self.calmar
        }
        if file_name is not None:
            self.res.to_excel(file_name)

    def plot_res(self) -> None:
        """
        画图，
        :return: None
        """
        fig, ax = plt.subplots(3, 1)
        ax[0].plot(self.date, self.res.loc[:, 'line_rnt'])
        ax[1].plot(self.date, self.inf)
        ax[2].plot(self.date, self.close)
        plt.show()

    def get_score(self, plot: bool = False, file_name: str = None,
                  score_type: str = 'shape') -> float:
        """
        获取评分
        :param plot: 是否画图
        :param file_name:文件名称，一般为None，只获取得分，不生成文件
        :param score_type: 评分标准，get_res结果字典中的项目，mean_shape_calmar： 夏普和calmar的均值
        :return: 最后的评分
        """
        res_d = self.get_res(plot, file_name)
        res_type = ['net', 'avg', 'std', 'shape', 'maxDD', 'calmar']
        if score_type in res_type:
            k = res_d[score_type]
        elif score_type == 'mean_shape_calmar':
            k = 0.5 * self.shape + 0.5 * self.calmar
        else:
            raise '评分类型不存在'
        return k

    def get_res(self, plot: bool = True, file_name: str = None) -> dict:
        """
        输出结果自字典，包括年化收益率，波动率， 夏普比率， 最大回撤， calmar比率等
        :param plot: 是否画图
        :param file_name: 文件名
        :return: 结果字典
        """
        res_type = ['net', 'avg', 'std', 'shape', 'maxDD', 'calmar']
        self.table(file_name=file_name)
        if plot:
            self.plot_res()
        d = {}
        for i in res_type:
            d[i] = self.back_res[i]
        return d
